import torch
import torchvision
from PIL import Image
from torch import nn
from torch.utils.data import DataLoader

test_data = torchvision.datasets.CIFAR10(root="./dataset",train=False,transform=torchvision.transforms.ToTensor(),download=True)
test_dataLoader = DataLoader(test_data, batch_size=64)

image_path = './imgs/dog.png'
# image_path = './imgs/cat.png'
image = Image.open(image_path).convert("RGB")

transform = torchvision.transforms.Compose([torchvision.transforms.Resize((32,32)),
                                            torchvision.transforms.ToTensor()])

image = transform(image)

class Tudui(nn.Module):
  def __init__(self):
    super(Tudui,self).__init__()
    self.model = nn.Sequential(
      nn.Conv2d(3, 32, 5, 1, 2),
      nn.MaxPool2d(2),
      nn.Conv2d(32, 32, 5, 1, 2),
      nn.MaxPool2d(2),
      nn.Conv2d(32, 64, 5, 1, 2),
      nn.MaxPool2d(2),
      nn.Flatten(),
      nn.Linear(1024, 64),
      nn.Linear(64, 10)
    )
  
  def forward(self, x):
    x = self.model(x)
    return x
  
model = torch.load('./saveModels/tudui_29_gpu.pth', map_location=torch.device('cpu'))
image = torch.reshape(image,(1,3,32,32))
model.eval()
with torch.no_grad():
  output = model(image)
  output_idx = output.argmax(1)
  output_label = test_data.classes[output_idx]
  print('识别预测的结果是：{}'.format(output_label))
